"""
    @Project: 聚类算法：k-means
    https://blog.csdn.net/xc_zhou/article/details/88247783
"""
import numpy as np
from matplotlib import pyplot


class K_Means(object):
    def __init__(self, k=2, tolerance=0.0001, max_iter=300):
        """
        :param k:分组数
        :param tolerance:中心点误差
        :param max_iter:迭代次数
        """
        self.k_ = k
        self.tolerance_ = tolerance
        self.max_iter_ = max_iter

    def fit(self, data):
        """
        k均值计算
        :param data:输入样本数据
        """
        self.centers_ = {}  # 中心点
        for i in range(self.k_):
            self.centers_[i] = data[i]

        for i in range(self.max_iter_):
            self.clf_ = {}  # 分组情况
            for i in range(self.k_):
                self.clf_[i] = []  # 每次迭代清空分组结果
            for feature in data:
                distances = []
                for center in self.centers_:
                    distances.append(np.linalg.norm(feature - self.centers_[center]))  # 欧式距离
                classification = distances.index(min(distances))  # 单个数据的分组结果
                self.clf_[classification].append(feature)  # 将单个数据添加到不同组中

            print("中心点:", self.centers_)
            print("分组情况:", self.clf_)
            print("----------------------------------------------------------------------")

            prev_centers = dict(self.centers_)
            for c in self.clf_:
                self.centers_[c] = np.average(self.clf_[c], axis=0)  # 重新计算中心点坐标

            # 中心点是否在误差范围
            optimized = True
            for center in self.centers_:
                org_centers = prev_centers[center]  # 上一次的中心点坐标
                cur_centers = self.centers_[center]  # 这一次的中心点坐标
                if np.sum((cur_centers - org_centers) / (org_centers * 100.0 + 1e-6)) > self.tolerance_:
                    optimized = False
            if optimized:
                break  # 两次中心点坐标比较相差无几后，结束循环

    def predict(self, p_data):
        """
        k均值预测数据
        :param p_data:输入预测数据
        :return:单个数据的分组结果
        """
        distances = [np.linalg.norm(p_data - self.centers_[center]) for center in self.centers_]  # 欧式距离
        index = distances.index(min(distances))   # 单个数据的分组结果
        return index


if __name__ == '__main__':
    x = np.array([[0, 0], [1, 0], [0, 1], [1, 1], [2, 1], [1, 2], [2, 2], [3, 2], [6, 6], [7, 6],
                  [8, 6], [6, 7], [7, 7], [8, 7], [9, 7], [7, 8], [8, 8], [9, 8], [8, 9], [9, 9]])
    k_means = K_Means(k=2)  # 分成2类
    k_means.fit(x)
    for center in k_means.centers_:
        pyplot.scatter(k_means.centers_[center][0], k_means.centers_[center][1], marker='*', s=150)  # 画出中心点

    for cat in k_means.clf_:
        for point in k_means.clf_[cat]:
            pyplot.scatter(point[0], point[1], c=('r' if cat == 0 else 'b'))  # 画出样本数据

    predict = [[2, 3], [6, 9]]
    for feature in predict:
        cat = k_means.predict(feature)
        pyplot.scatter(feature[0], feature[1], c=('r' if cat == 0 else 'b'), marker='x')  # 画出预测数据

    pyplot.show()